use anyhow::Result;
use clap::Parser;
use futures_lite::StreamExt;
use iroh::{Endpoint, NodeAddr, NodeId, protocol::Router};
use iroh_gossip::{
    net::{Event, Gossip, GossipEvent, GossipReceiver},
    proto::TopicId,
};
use rusqlite::ffi::sqlite3_auto_extension;
use serde::{Deserialize, Serialize};
use sqlite_vec::sqlite3_vec_init;
use std::{collections::HashMap, fmt, str::FromStr};
mod embed;
mod ai;
use ai::AIAgent;

/// Chat over iroh-gossip
///
/// This broadcasts unsigned messages over iroh-gossip.
///
/// By default a new node id is created when starting the example.
///
/// By default, we use the default n0 discovery services to dial by `NodeId`.
#[derive(Parser, Debug)]
struct Args {
    /// Set your nickname.
    #[clap(short, long)]
    name: Option<String>,
    /// Set the bind port for our socket. By default, a random port will be used.
    #[clap(short, long, default_value = "0")]
    bind_port: u16,
    #[clap(subcommand)]
    command: Command,
}

#[derive(Parser, Debug)]
enum Command {
    /// Open a chat room for a topic and print a ticket for others to join.
    Open,
    /// Join a chat room from a ticket.
    Join {
        /// The ticket, as base32 string.
        ticket: String,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    unsafe {
        sqlite3_auto_extension(Some(std::mem::transmute(sqlite3_vec_init as *const ())));
    }

    let args = Args::parse();

    // parse the cli command
    let (topic, nodes) = match &args.command {
        Command::Open => {
            let topic = TopicId::from_bytes(rand::random());
            println!("> opening chat room for topic {topic}");
            (topic, vec![])
        }
        Command::Join { ticket } => {
            let Ticket { topic, nodes } = Ticket::from_str(ticket)?;
            println!("> joining chat room for topic {topic}");
            (topic, nodes)
        }
    };

    let endpoint = Endpoint::builder().discovery_n0().bind().await?;

    println!("> our node id: {}", endpoint.node_id());
    let gossip = Gossip::builder().spawn(endpoint.clone()).await?;

    let router = Router::builder(endpoint.clone())
        .accept(iroh_gossip::ALPN, gossip.clone())
        .spawn();

    // in our main file, after we create a topic `id`:
    // print a ticket that includes our own node id and endpoint addresses
    let ticket = {
        // Get our address information, includes our
        // `NodeId`, our `RelayUrl`, and any direct
        // addresses.
        let me = endpoint.node_addr().await?;
        let nodes = vec![me];
        Ticket { topic, nodes }
    };
    println!("> ticket to join us: {ticket}");

    // join the gossip topic by connecting to known nodes, if any
    let node_ids = nodes.iter().map(|p| p.node_id).collect();
    if nodes.is_empty() {
        println!("> waiting for nodes to join us...");
    } else {
        println!("> trying to connect to {} nodes...", nodes.len());
        // add the peer addrs from the ticket to our endpoint's addressbook so that they can be dialed
        for node in nodes.into_iter() {
            endpoint.add_node_addr(node)?;
        }
    };
    let (sender, receiver) = gossip.subscribe_and_join(topic, node_ids).await?.split();
    println!("> connected!");

    // 提取用户名以便复用
    let my_name = args.name.clone();

    // broadcast our name, if set
    if let Some(name) = args.name {
        let message = Message::new(MessageBody::AboutMe {
            from: endpoint.node_id(),
            name,
        });
        sender.broadcast(message.to_vec().into()).await?;
    }

    // subscribe and print loop
    let my_node_id = endpoint.node_id();
    tokio::spawn(subscribe_loop(receiver, my_node_id, my_name.clone()));

    // spawn an input thread that reads stdin
    // create a multi-provider, single-consumer channel
    let (line_tx, mut line_rx) = tokio::sync::mpsc::channel(1);
    // and pass the `sender` portion to the `input_loop`
    std::thread::spawn(move || input_loop(line_tx));

    // broadcast each line we type
    println!("> type a message and hit enter to broadcast...");
    println!("> you can use @username to mention someone");
    println!("> you can use /username message to send private message");
    println!("> you can use @ai to ask questions to the AI assistant");

    // 初始化 AI agent (异步)
    let ai_agent = match AIAgent::new().await {
        Ok(agent) => {
            println!("🤖 AI assistant is ready!");
            Some(agent)
        }
        Err(_) => {
            println!(
                "⚠️  AI assistant is not available. Set DEEPSEEK_API_KEY to enable AI features."
            );
            None
        }
    };

    // listen for lines that we have typed to be sent from `stdin`
    while let Some(text) = line_rx.recv().await {
        let text = text.trim().to_string();

        // 首先检查是否是私人消息
        let message = if let Some((target_user, private_text)) = parse_private_message(&text) {
            // 创建私人消息
            Message::new(MessageBody::PrivateMessage {
                from: endpoint.node_id(),
                to: target_user.clone(),
                text: private_text.clone(),
            })
        } else {
            // 解析消息中的 @mentions
            let mentions = parse_mentions(&text);

            // （@ai 提及检查移到后面处理）

            // 根据是否包含提及创建不同类型的消息
            if mentions.is_empty() {
                Message::new(MessageBody::Message {
                    from: endpoint.node_id(),
                    text: text.clone(),
                })
            } else {
                Message::new(MessageBody::MentionMessage {
                    from: endpoint.node_id(),
                    text: text.clone(),
                    mentions,
                })
            }
        };

        // broadcast the encoded message
        sender.broadcast(message.to_vec().into()).await?;

        // print to ourselves the text that we sent
        match &message.body {
            MessageBody::PrivateMessage { to, text, .. } => {
                println!("> sent private message to {}: {}", to, text);
            }
            MessageBody::MentionMessage { text, .. } => {
                println!("> sent: {}", text);

                // 处理 @ai 提及
                if AIAgent::is_ai_mention(text) && ai_agent.is_some() {
                    if let Some(user_name) = &my_name {
                        println!("🤖 Processing AI request...");

                        // 异步处理 AI 请求
                        let question = AIAgent::extract_question(text);
                        let sender_clone = sender.clone();
                        let node_id = endpoint.node_id();
                        let user_name_clone = user_name.clone();
                        let question_clone = question.clone();

                        let ai_agent_clone = ai_agent.clone();
                        tokio::spawn(async move {
                            if let Some(agent) = ai_agent_clone {
                                match agent.process_question(&question, &user_name_clone).await {
                                    Ok(answer) => {
                                        let ai_response = Message::new(MessageBody::AIResponse {
                                            from: node_id,
                                            to: user_name_clone.to_string(),
                                            question: question_clone,
                                            answer,
                                        });

                                        if let Err(e) = sender_clone
                                            .broadcast(ai_response.to_vec().into())
                                            .await
                                        {
                                            println!("❌ Failed to send AI response: {}", e);
                                        } else {
                                            println!("🤖 AI response sent");
                                        }
                                    }
                                    Err(e) => {
                                        println!("❌ AI processing error: {}", e);
                                    }
                                }
                            }
                        });
                    } else {
                        println!("⚠️  Cannot process AI request: no username set");
                    }
                }
            }
            _ => {
                println!("> sent: {text}");
            }
        }
    }
    router.shutdown().await?;

    Ok(())
}

#[derive(Debug, Serialize, Deserialize)]
struct Message {
    body: MessageBody,
    nonce: [u8; 16],
}

#[derive(Debug, Serialize, Deserialize)]
enum MessageBody {
    AboutMe {
        from: NodeId,
        name: String,
    },
    Message {
        from: NodeId,
        text: String,
    },
    MentionMessage {
        from: NodeId,
        text: String,
        mentions: Vec<String>, // 被提及的用户名列表
    },
    PrivateMessage {
        from: NodeId,
        to: String, // 目标用户名
        text: String,
    },
    AIResponse {
        from: NodeId,
        to: String,       // 原始提问者的用户名
        question: String, // 原始问题
        answer: String,   // AI 回答
    },
}

impl Message {
    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        serde_json::from_slice(bytes).map_err(Into::into)
    }

    pub fn new(body: MessageBody) -> Self {
        Self {
            body,
            nonce: rand::random(),
        }
    }

    pub fn to_vec(&self) -> Vec<u8> {
        serde_json::to_vec(self).expect("serde_json::to_vec is infallible")
    }
}

// Handle incoming events
async fn subscribe_loop(
    mut receiver: GossipReceiver,
    my_node_id: NodeId,
    my_name: Option<String>,
) -> Result<()> {
    // keep track of the mapping between `NodeId`s and names
    let mut names = HashMap::new();
    // 反向映射：从名字到NodeId
    let mut name_to_node = HashMap::new();

    // 如果我们有自己的昵称，先添加到映射中
    if let Some(name) = &my_name {
        names.insert(my_node_id, name.clone());
        name_to_node.insert(name.clone(), my_node_id);
    }

    // iterate over all events
    while let Some(event) = receiver.try_next().await? {
        // if the Event is a `GossipEvent::Received`, let's deserialize the message:
        if let Event::Gossip(GossipEvent::Received(msg)) = event {
            // deserialize the message and match on the
            // message type:
            match Message::from_bytes(&msg.content)?.body {
                MessageBody::AboutMe { from, name } => {
                    // if it's an `AboutMe` message
                    // add and entry into the map
                    // and print the name
                    names.insert(from, name.clone());
                    name_to_node.insert(name.clone(), from);
                    println!("> {} is now known as {}", from.fmt_short(), name);
                }
                MessageBody::Message { from, text } => {
                    // if it's a `Message` message,
                    // get the name from the map
                    // and print the message
                    let name = names
                        .get(&from)
                        .map_or_else(|| from.fmt_short(), String::to_string);
                    println!("{}: {}", name, text);
                }
                MessageBody::MentionMessage {
                    from,
                    text,
                    mentions,
                } => {
                    // 处理包含提及的消息
                    let sender_name = names
                        .get(&from)
                        .map_or_else(|| from.fmt_short(), String::to_string);

                    // 检查当前用户是否被提及
                    let my_name = names.get(&my_node_id);
                    let is_mentioned = mentions.iter().any(|mention| {
                        // 检查是否通过用户名被提及
                        if let Some(my_name) = my_name {
                            mention == my_name
                        } else {
                            false
                        }
                    });

                    if is_mentioned {
                        println!("🔔 {} mentioned you: {}", sender_name, text);
                    } else {
                        println!("{}: {}", sender_name, text);
                    }

                    // 显示所有被提及的用户
                    if !mentions.is_empty() {
                        println!("   └─ mentions: {}", mentions.join(", "));
                    }
                }
                MessageBody::PrivateMessage { from, to, text } => {
                    // 处理私人消息 - 只有目标用户才能看到
                    let current_user_name = names.get(&my_node_id);

                    // 检查消息是否是给当前用户的
                    let is_for_me = if let Some(my_name_in_map) = current_user_name {
                        &to == my_name_in_map
                    } else if let Some(my_actual_name) = &my_name {
                        // 如果names映射中没有我们的名字，检查全局的my_name
                        &to == my_actual_name
                    } else {
                        false
                    };

                    if is_for_me {
                        let sender_name = names
                            .get(&from)
                            .map_or_else(|| from.fmt_short(), String::to_string);
                        println!("💬 [Private] {}: {}", sender_name, text);
                    }
                    // 如果不是给我的，就静默忽略这条消息
                }
                MessageBody::AIResponse {
                    from: _,
                    to,
                    question,
                    answer,
                } => {
                    // 处理 AI 回复 - 只有原始提问者才能看到
                    let current_user_name = names.get(&my_node_id);

                    // 检查是否是给当前用户的 AI 回复
                    let is_for_me = if let Some(my_name_in_map) = current_user_name {
                        &to == my_name_in_map
                    } else if let Some(my_actual_name) = &my_name {
                        &to == my_actual_name
                    } else {
                        false
                    };

                    if is_for_me {
                        println!("🤖 [AI Response] Re: \"{}\"", question);
                        println!("   {}", answer);
                    }
                    // 如果不是给我的，就静默忽略这条消息
                }
            }
        }
    }
    Ok(())
}

fn input_loop(line_tx: tokio::sync::mpsc::Sender<String>) -> Result<()> {
    let mut buffer = String::new();
    let stdin = std::io::stdin(); // We get `Stdin` here.
    loop {
        stdin.read_line(&mut buffer)?;
        line_tx.blocking_send(buffer.clone())?;
        buffer.clear();
    }
}

/// 解析消息中的 @mentions
fn parse_mentions(text: &str) -> Vec<String> {
    let mut mentions = Vec::new();
    let words: Vec<&str> = text.split_whitespace().collect();

    for word in words {
        if word.starts_with('@') && word.len() > 1 {
            let username = &word[1..]; // 去掉 @ 符号
            // 移除可能的标点符号
            let clean_username =
                username.trim_end_matches(&[',', '.', '!', '?', ':', ';', ')', ']', '}'][..]);
            if !clean_username.is_empty() {
                mentions.push(clean_username.to_string());
            }
        }
    }

    // 去重
    mentions.sort();
    mentions.dedup();
    mentions
}

/// 解析私人消息格式：/username message
/// 返回：(目标用户名, 消息内容)，如果不是私人消息格式则返回 None
fn parse_private_message(text: &str) -> Option<(String, String)> {
    let text = text.trim();
    if text.starts_with('/') && text.len() > 1 {
        // 找到第一个空格的位置
        if let Some(space_pos) = text.find(' ') {
            let username = &text[1..space_pos]; // 去掉 / 符号
            let message = text[space_pos + 1..].trim();

            if !username.is_empty() && !message.is_empty() {
                return Some((username.to_string(), message.to_string()));
            }
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_mentions() {
        // 测试基本提及
        assert_eq!(parse_mentions("Hello @alice"), vec!["alice"]);

        // 测试多个提及
        assert_eq!(
            parse_mentions("@bob @charlie hello"),
            vec!["bob", "charlie"]
        );

        // 测试带标点符号的提及
        assert_eq!(parse_mentions("Thanks @alice!"), vec!["alice"]);
        assert_eq!(parse_mentions("Hey @bob, how are you?"), vec!["bob"]);

        // 测试重复提及
        assert_eq!(parse_mentions("@alice @alice"), vec!["alice"]);

        // 测试空字符串和无提及
        assert_eq!(parse_mentions(""), Vec::<String>::new());
        assert_eq!(parse_mentions("Hello world"), Vec::<String>::new());

        // 测试单独的@符号
        assert_eq!(parse_mentions("@"), Vec::<String>::new());
    }

    #[test]
    fn test_parse_private_message() {
        // 测试基本私人消息
        assert_eq!(
            parse_private_message("/alice Hello there!"),
            Some(("alice".to_string(), "Hello there!".to_string()))
        );

        // 测试带空格的消息
        assert_eq!(
            parse_private_message("/bob How are you doing today?"),
            Some(("bob".to_string(), "How are you doing today?".to_string()))
        );

        // 测试前后有空格的情况
        assert_eq!(
            parse_private_message("  /charlie  Hi!  "),
            Some(("charlie".to_string(), "Hi!".to_string()))
        );

        // 测试非私人消息格式
        assert_eq!(parse_private_message("Hello world"), None);
        assert_eq!(parse_private_message("@alice hello"), None);
        assert_eq!(parse_private_message(""), None);

        // 测试无效格式
        assert_eq!(parse_private_message("/"), None);
        assert_eq!(parse_private_message("/alice"), None); // 没有消息内容
        assert_eq!(parse_private_message("/ hello"), None); // 没有用户名
    }
}

// add the `Ticket` code to the bottom of the main file
#[derive(Debug, Serialize, Deserialize)]
struct Ticket {
    topic: TopicId,
    nodes: Vec<NodeAddr>,
}

impl Ticket {
    /// Deserialize from a slice of bytes to a Ticket.
    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        serde_json::from_slice(bytes).map_err(Into::into)
    }

    /// Serialize from a `Ticket` to a `Vec` of bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        serde_json::to_vec(self).expect("serde_json::to_vec is infallible")
    }
}

// The `Display` trait allows us to use the `to_string`
// method on `Ticket`.
impl fmt::Display for Ticket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut text = data_encoding::BASE32_NOPAD.encode(&self.to_bytes()[..]);
        text.make_ascii_lowercase();
        write!(f, "{}", text)
    }
}

// The `FromStr` trait allows us to turn a `str` into
// a `Ticket`
impl FromStr for Ticket {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let bytes = data_encoding::BASE32_NOPAD.decode(s.to_ascii_uppercase().as_bytes())?;
        Self::from_bytes(&bytes)
    }
}
